# Variable pool system for FLEXINVERT MCMC
# Inspired by MCIntegration.jl's adaptive variable management

module Variables

using LinearAlgebra
using Statistics
using Random

export AbstractFluxVariable, ContinuousFlux, DiscreteFlux, CompositeFlux
export create!, remove!, shift!, swap!, train!, reset_statistics!
export shiftRollback!, swapRollback!, removeRollback!, createRollback!

"""
    AbstractFluxVariable{T}

Abstract base type for FLEXINVERT flux variables with adaptive sampling.
All variable types must implement the core MCMC operations.
"""
abstract type AbstractFluxVariable{T} <: AbstractVector{T} end

"""
    VegasGrid

Adaptive importance sampling grid for continuous variables.
Similar to MCIntegration.jl's grid system but specialized for flux parameters.
"""
mutable struct VegasGrid
    grid::Vector{Float64}           # Grid boundaries
    histogram::Vector{Float64}      # Accumulated histogram for adaptation
    n_bins::Int                     # Number of grid bins
    alpha::Float64                  # Learning rate for adaptation
    adapt::Bool                     # Enable/disable adaptation

    function VegasGrid(n_bins::Int = 100, alpha::Float64 = 0.5)
        grid = collect(range(0.0, 1.0, length=n_bins+1))
        histogram = zeros(Float64, n_bins)
        new(grid, histogram, n_bins, alpha, true)
    end
end

"""
    ContinuousFlux{T} <: AbstractFluxVariable{T}

Continuous flux variables with Vegas-style adaptive importance sampling.
Represents physical flux parameters that can vary continuously.
"""
mutable struct ContinuousFlux{T} <: AbstractFluxVariable{T}
    data::Vector{T}                 # Variable values in [0,1] space
    physical::Vector{T}             # Values in physical flux space
    bounds::Vector{Tuple{T,T}}      # Physical bounds for each variable
    grid::VegasGrid                 # Adaptive sampling grid
    prob::Vector{Float64}           # Current proposal probabilities

    # Statistics tracking
    n_proposed::Vector{Int}         # Proposals per variable
    n_accepted::Vector{Int}         # Acceptances per variable

    # Backup for rollback
    backup_data::Vector{T}
    backup_physical::Vector{T}
    backup_prob::Vector{Float64}

    function ContinuousFlux{T}(n_vars::Int, bounds::Vector{Tuple{T,T}}) where T
        data = Vector{T}(undef, n_vars)
        physical = Vector{T}(undef, n_vars)
        grid = VegasGrid()
        prob = ones(Float64, n_vars)

        new{T}(
            data, physical, bounds, grid, prob,
            zeros(Int, n_vars), zeros(Int, n_vars),
            Vector{T}(undef, n_vars), Vector{T}(undef, n_vars),
            Vector{Float64}(undef, n_vars)
        )
    end
end

ContinuousFlux(n_vars::Int, bounds::Vector{Tuple{T,T}}) where T = ContinuousFlux{T}(n_vars, bounds)

# Vector interface implementation
Base.size(var::ContinuousFlux) = size(var.data)
Base.getindex(var::ContinuousFlux, i::Int) = var.data[i]
Base.setindex!(var::ContinuousFlux, v, i::Int) = (var.data[i] = v; update_physical!(var, i))

"""
    update_physical!(var::ContinuousFlux, idx::Int)

Update physical space representation from [0,1] space using adaptive grid.
"""
function update_physical!(var::ContinuousFlux{T}, idx::Int) where T
    u = var.data[idx]
    bounds = var.bounds[idx]

    # Map through Vegas grid
    bin_idx = max(1, min(var.grid.n_bins, Int(ceil(u * var.grid.n_bins))))
    local_u = (u * var.grid.n_bins) - (bin_idx - 1)

    # Linear interpolation within bin
    grid_val = var.grid.grid[bin_idx] + local_u * (var.grid.grid[bin_idx+1] - var.grid.grid[bin_idx])

    # Map to physical bounds
    var.physical[idx] = bounds[1] + grid_val * (bounds[2] - bounds[1])
end

"""
    create!(var::ContinuousFlux, idx::Int, config)

Create a new variable at the specified index.
Returns the proposal probability factor.
"""
function create!(var::ContinuousFlux{T}, idx::Int, config) where T
    # Generate uniform random value in [0,1]
    var.data[idx] = rand(config.rng, T)
    update_physical!(var, idx)

    # Proposal probability is uniform
    var.prob[idx] = 1.0
    return 1.0
end

"""
    remove!(var::ContinuousFlux, idx::Int, config)

Remove a variable at the specified index.
Returns the proposal probability factor.
"""
function remove!(var::ContinuousFlux, idx::Int, config)
    # For removal, probability factor is inverse of current density
    prob_factor = 1.0 / var.prob[idx]

    # Clear the variable (set to zero for removal)
    var.data[idx] = zero(eltype(var.data))
    var.physical[idx] = zero(eltype(var.physical))
    var.prob[idx] = 0.0

    return prob_factor
end

"""
    shift!(var::ContinuousFlux, idx::Int, config)

Propose a new value for an existing variable.
Returns the proposal probability factor.
"""
function shift!(var::ContinuousFlux{T}, idx::Int, config) where T
    # Store current values for rollback
    var.backup_data[idx] = var.data[idx]
    var.backup_physical[idx] = var.physical[idx]
    var.backup_prob[idx] = var.prob[idx]

    # Gaussian random walk in [0,1] space
    step_size = get(config, :flux_step_size, 0.1)
    new_val = var.data[idx] + step_size * randn(config.rng, T)

    # Reflect at boundaries to stay in [0,1]
    while new_val < 0 || new_val > 1
        if new_val < 0
            new_val = -new_val
        elseif new_val > 1
            new_val = 2 - new_val
        end
    end

    # Update variable
    var.data[idx] = new_val
    update_physical!(var, idx)

    # For symmetric proposals, probability factor is 1
    # (Vegas adaptation affects this, but we'll keep it simple for now)
    var.prob[idx] = 1.0
    var.n_proposed[idx] += 1

    return 1.0
end

"""
    swap!(var::ContinuousFlux, idx1::Int, idx2::Int, config)

Swap two variables.
Returns the proposal probability factor.
"""
function swap!(var::ContinuousFlux, idx1::Int, idx2::Int, config)
    if idx1 == idx2
        return 1.0
    end

    # Store for rollback
    var.backup_data[idx1] = var.data[idx1]
    var.backup_data[idx2] = var.data[idx2]
    var.backup_physical[idx1] = var.physical[idx1]
    var.backup_physical[idx2] = var.physical[idx2]
    var.backup_prob[idx1] = var.prob[idx1]
    var.backup_prob[idx2] = var.prob[idx2]

    # Perform swap
    var.data[idx1], var.data[idx2] = var.data[idx2], var.data[idx1]
    var.physical[idx1], var.physical[idx2] = var.physical[idx2], var.physical[idx1]
    var.prob[idx1], var.prob[idx2] = var.prob[idx2], var.prob[idx1]

    # For symmetric swap, probability factor is 1
    return 1.0
end

# Rollback functions
"""
    shiftRollback!(var::ContinuousFlux, idx::Int, config)

Rollback a shift operation.
"""
function shiftRollback!(var::ContinuousFlux, idx::Int, config)
    var.data[idx] = var.backup_data[idx]
    var.physical[idx] = var.backup_physical[idx]
    var.prob[idx] = var.backup_prob[idx]
end

"""
    swapRollback!(var::ContinuousFlux, idx1::Int, idx2::Int, config)

Rollback a swap operation.
"""
function swapRollback!(var::ContinuousFlux, idx1::Int, idx2::Int, config)
    if idx1 == idx2
        return
    end

    var.data[idx1] = var.backup_data[idx1]
    var.data[idx2] = var.backup_data[idx2]
    var.physical[idx1] = var.backup_physical[idx1]
    var.physical[idx2] = var.backup_physical[idx2]
    var.prob[idx1] = var.backup_prob[idx1]
    var.prob[idx2] = var.backup_prob[idx2]
end

"""
    removeRollback!(var::ContinuousFlux, idx::Int, config)

Rollback a remove operation.
"""
function removeRollback!(var::ContinuousFlux, idx::Int, config)
    # For remove rollback, we need to restore the previous state
    # This assumes backup was set during a previous operation
    var.data[idx] = var.backup_data[idx]
    var.physical[idx] = var.backup_physical[idx]
    var.prob[idx] = var.backup_prob[idx]
end

"""
    createRollback!(var::ContinuousFlux, idx::Int, config)

Rollback a create operation.
"""
function createRollback!(var::ContinuousFlux, idx::Int, config)
    # For create rollback, remove the created variable
    var.data[idx] = zero(eltype(var.data))
    var.physical[idx] = zero(eltype(var.physical))
    var.prob[idx] = 0.0
end

"""
    train!(var::ContinuousFlux)

Train the adaptive Vegas grid based on accumulated histogram.
"""
function train!(var::ContinuousFlux)
    if !var.grid.adapt || sum(var.grid.histogram) == 0
        return
    end

    # Smooth the histogram
    smoothed = smooth_histogram(var.grid.histogram, 6.0)

    # Rescale for numerical stability
    smoothed = rescale_histogram(smoothed, var.grid.alpha)

    # Redistribute grid points
    redistribute_grid!(var.grid, smoothed)

    # Reset histogram for next adaptation cycle
    fill!(var.grid.histogram, 0.0)
end

"""
    smooth_histogram(hist::Vector{Float64}, factor::Float64)

Smooth histogram to reduce noise in adaptation.
"""
function smooth_histogram(hist::Vector{Float64}, factor::Float64)
    n = length(hist)
    smoothed = copy(hist)

    for i in 2:n-1
        smoothed[i] = (hist[i-1] + factor*hist[i] + hist[i+1]) / (factor + 2.0)
    end

    return smoothed
end

"""
    rescale_histogram(hist::Vector{Float64}, alpha::Float64)

Rescale histogram for numerical stability.
"""
function rescale_histogram(hist::Vector{Float64}, alpha::Float64)
    # Avoid too extreme values
    log_hist = log.(max.(hist, 1e-10))
    log_avg = mean(log_hist)

    # Rescale towards average
    log_rescaled = log_avg .+ alpha .* (log_hist .- log_avg)

    return exp.(log_rescaled)
end

"""
    redistribute_grid!(grid::VegasGrid, distribution::Vector{Float64})

Redistribute grid points based on learned importance distribution.
"""
function redistribute_grid!(grid::VegasGrid, distribution::Vector{Float64})
    n_bins = length(distribution)
    total_weight = sum(distribution)

    if total_weight <= 0
        return
    end

    new_grid = similar(grid.grid)
    new_grid[1] = 0.0
    new_grid[end] = 1.0

    # Cumulative distribution
    cumsum_dist = cumsum(distribution) / total_weight

    # Redistribute grid points
    for i in 2:(length(new_grid)-1)
        target_weight = (i-1) / (length(new_grid)-1)

        # Find corresponding position in old grid
        bin_idx = searchsortedfirst(cumsum_dist, target_weight)
        bin_idx = max(1, min(bin_idx, n_bins))

        if bin_idx == 1
            new_grid[i] = target_weight * (grid.grid[2] - grid.grid[1])
        else
            # Linear interpolation
            fraction = (target_weight - cumsum_dist[bin_idx-1]) /
                      (cumsum_dist[bin_idx] - cumsum_dist[bin_idx-1])
            new_grid[i] = grid.grid[bin_idx] + fraction * (grid.grid[bin_idx+1] - grid.grid[bin_idx])
        end
    end

    grid.grid .= new_grid
end

"""
    reset_statistics!(var::ContinuousFlux)

Reset adaptation statistics.
"""
function reset_statistics!(var::ContinuousFlux)
    fill!(var.n_proposed, 0)
    fill!(var.n_accepted, 0)
    fill!(var.grid.histogram, 0.0)
end

"""
    acceptance_rate(var::ContinuousFlux, idx::Int)

Get acceptance rate for a specific variable.
"""
function acceptance_rate(var::ContinuousFlux, idx::Int)
    if var.n_proposed[idx] == 0
        return 0.0
    end
    return var.n_accepted[idx] / var.n_proposed[idx]
end

"""
    accept!(var::ContinuousFlux, idx::Int)

Mark a proposal as accepted for adaptation.
"""
function accept!(var::ContinuousFlux, idx::Int)
    var.n_accepted[idx] += 1

    # Update histogram for Vegas adaptation
    u = var.data[idx]
    bin_idx = max(1, min(var.grid.n_bins, Int(ceil(u * var.grid.n_bins))))
    var.grid.histogram[bin_idx] += 1.0
end

end # module Variables